# encoding :utf-8

import io  # 文件数据流
import datetime
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow import keras
# 导入常见网络层, sequential容器, 优化器, 损失函数
from tensorflow.keras import layers, Sequential, optimizers, losses, metrics, utils
import os # 运维模块， 调用系统命令
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 只显示warring和error


def preprocess(x, y):
    x = tf.cast(x, dtype=tf.float32) / 255.
    x = tf.reshape(x, (-1, 28 * 28))
    y = tf.cast(y, dtype=tf.int32)
    return x, y


batchsz = 128
(x, y), (x_val, y_val) = tf.keras.datasets.mnist.load_data()

# 训练数据
# 将数据存储到数据管道中，提速
db = tf.data.Dataset.from_tensor_slices((x,y))
# 特征归一化，标签转成整数形式
db = db.map(preprocess).shuffle(60000).batch(batchsz)
# 验证数据
ds_val = tf.data.Dataset.from_tensor_slices((x_val, y_val))
ds_val = ds_val.map(preprocess).batch(batchsz, drop_remainder=True) # drop_remainder 清除不足一批次的数据

network = Sequential([
    layers.Dense(256, activation='relu'),
    layers.Dense(128, activation='relu'),
    layers.Dense(64, activation='relu'),
    layers.Dense(32, activation='relu'),
    layers.Dense(10, activation='softmax')
])

# build  输入数据的维度
network.build(input_shape=(None, 28*28))
network.summary()

#定义log_dir:注意keras转义格式
log_dir=r"logs\\fix\\" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
# 数据回调到tensorboard中
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)

network.compile(optimizer=tf.keras.optimizers.Adam(lr=0.01),
              loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])

network.fit(db, epochs=5,callbacks=[tensorboard_callback])


